import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors

# --- 1. 全局样式和准备工作 ---
sns.set_theme(style="whitegrid")
# 文件路径和组名保持不变
# 注意：请确保这些路径在您的环境中是正确的
file_paths = [
    '/home/liudd/deeplearing/droupout/result/train2.1_mc_prediction_results.csv',
    '/home/liudd/deeplearing/droupout/result/train3.1_mc_prediction_results.csv',
    '/home/liudd/deeplearing/droupout/result/train4.1_mc_prediction_results.csv',
    '/home/liudd/deeplearing/droupout/result/train5.1_mc_prediction_results.csv',
    '/home/liudd/deeplearing/droupout/result/train6.1_mc_prediction_results.csv'
]
group_names = ['water', 'supercooled', 'mix', 'ice', 'cirrus']

# --- 2. 定义通用的分箱规则 ---
cbh_bins = [0, 500, 1000, 2000, 4000, 8000, 20000]
cbh_labels = ['0-0.5km', '0.5-1km', '1-2km', '2-4km', '4-8km', '>8km']

thickness_bins = [0, 500, 1000, 2000, 4000, 8000, 20000]
thickness_labels = ['0-0.5km', '0.5-1km', '1-2km', '2-4km', '4-8km', '>8km']

# --- 3. 数据预加载和处理 (一次性完成) ---
print("正在加载和预处理数据...")
all_data_loaded = []
for filepath in file_paths:
    try:
        df = pd.read_csv(filepath)
        df['y_val_true'] = df['cloudsat_cbh'].values
        df['y_pred_true'] = df['Predicted_Mean_CBH'].values
        # 关键步骤：只有在分母不为0时才计算相对误差
        df['rel_err'] = np.abs(df['y_val_true'] - df['y_pred_true']) / np.where(df['y_val_true'] == 0, np.nan, df['y_val_true'])
        df['cloud_thickness'] = df['cloudsat_cth'] - df['cloudsat_cbh']

        df.replace([np.inf, -np.inf], np.nan, inplace=True)
        df.dropna(subset=['rel_err', 'cloudsat_cbh', 'cloud_thickness'], inplace=True)

        # 进行分箱
        df['cbh_bin'] = pd.cut(df['cloudsat_cbh'], bins=cbh_bins, labels=cbh_labels, right=False)
        df['thickness_bin'] = pd.cut(df['cloud_thickness'], bins=thickness_bins, labels=thickness_labels, right=False)
        all_data_loaded.append(df)
    except FileNotFoundError:
        print(f"警告: 文件未找到 {filepath}, 将跳过。")
        # 添加一个空的dataframe以保持索引一致
        all_data_loaded.append(pd.DataFrame())


# 合并所有数据
combined_df_all_types = pd.concat([df for df in all_data_loaded if not df.empty], ignore_index=True)
print("数据处理完成。")


# ==============================================================================
# 图 1: 平均相对误差 (Mean Relative Error)
# ==============================================================================
print("\n正在生成图 1: 平均相对误差热力图...")
fig1, axes1 = plt.subplots(2, 3, figsize=(18, 10), sharex=True, sharey=True)
axes1 = axes1.flatten()

# 循环绘制每个云类型
for i, df in enumerate(all_data_loaded):
    if df.empty: continue
    ax = axes1[i]
    heatmap_data_err = df.groupby(['cbh_bin', 'thickness_bin'], observed=False)['rel_err'].mean().unstack()
    sns.heatmap(heatmap_data_err, annot=True, fmt=".2f", cmap="YlOrRd", linewidths=.5, ax=ax, cbar=False, vmin=0, vmax=2.5)
    ax.set_title(f'{group_names[i]}', fontsize=14, fontweight='bold')
    ax.set_ylabel('CloudSat CBH')
    ax.set_xlabel('Cloud Thickness')
    ax.invert_yaxis()

# 绘制 "All Types"
if not combined_df_all_types.empty:
    ax_all_err = axes1[-1]
    combined_heatmap_data_err = combined_df_all_types.groupby(['cbh_bin', 'thickness_bin'], observed=False)['rel_err'].mean().unstack()
    sns.heatmap(combined_heatmap_data_err, annot=True, fmt=".2f", cmap="YlOrRd", linewidths=.5, ax=ax_all_err, cbar=False, vmin=0, vmax=2.5)
    ax_all_err.set_title('All Types', fontsize=14, fontweight='bold')
    ax_all_err.set_ylabel('CloudSat CBH Bin')
    ax_all_err.set_xlabel('Cloud Thickness Bin')
    ax_all_err.invert_yaxis()

# 添加颜色条
cbar_ax_err = fig1.add_axes([0.92, 0.15, 0.02, 0.7])
norm_err = plt.Normalize(vmin=0, vmax=2.5)
fig1.colorbar(plt.cm.ScalarMappable(cmap="YlOrRd", norm=norm_err), cax=cbar_ax_err, label='Mean Relative Error')
# fig1.suptitle('图 1: 各云类型下模型预测的平均相对误差', fontsize=20)
# fig1.tight_layout(rect=[0, 0, 0.9, 0.95])
plt.show()


# ==============================================================================
# 图 2: 样本数量概率密度 (Sample Count Probability Density)
# ==============================================================================
print("\n正在生成图 2: 样本数量概率密度热力图...")
fig2, axes2 = plt.subplots(2, 3, figsize=(18, 10), sharex=True, sharey=True)
axes2 = axes2.flatten()

for i, df in enumerate(all_data_loaded):
    if df.empty: continue
    ax = axes2[i]
    sample_counts = df.groupby(['cbh_bin', 'thickness_bin'], observed=False).size().unstack(fill_value=0)
    total_samples = sample_counts.sum().sum()
    probability_density = sample_counts / total_samples if total_samples > 0 else sample_counts
    sns.heatmap(probability_density, annot=True, fmt=".3f", cmap="Blues", linewidths=.5, ax=ax, cbar=False, vmin=0, vmax=0.5)
    ax.set_title(f'{group_names[i]}', fontsize=14, fontweight='bold')
    ax.set_ylabel('CloudSat CBH')
    ax.set_xlabel('Cloud Thickness')
    ax.invert_yaxis()

# 绘制 "All Types"
if not combined_df_all_types.empty:
    ax_all_prob = axes2[-1]
    combined_counts = combined_df_all_types.groupby(['cbh_bin', 'thickness_bin'], observed=False).size().unstack(fill_value=0)
    combined_prob = combined_counts / combined_counts.sum().sum()
    sns.heatmap(combined_prob, annot=True, fmt=".3f", cmap="Blues", linewidths=.5, ax=ax_all_prob, cbar=False, vmin=0, vmax=0.5)
    ax_all_prob.set_title('All Types', fontsize=14, fontweight='bold')
    ax_all_prob.set_ylabel('CloudSat CBH Bin')
    ax_all_prob.set_xlabel('Cloud Thickness Bin')
    ax_all_prob.invert_yaxis()

# 添加颜色条
cbar_ax_prob = fig2.add_axes([0.92, 0.15, 0.02, 0.7])
norm_prob = plt.Normalize(vmin=0, vmax=0.5)
fig2.colorbar(plt.cm.ScalarMappable(cmap="Blues", norm=norm_prob), cax=cbar_ax_prob, label='Sample Count Probability Density')
# fig2.suptitle('图 2: 各云类型下样本数量的概率密度分布', fontsize=20)
# fig2.tight_layout(rect=[0, 0, 0.9, 0.95])
plt.show()


# ==============================================================================
# 图 3: 原始样本数量 (Raw Sample Count)
# ==============================================================================
print("\n正在生成图 3: 原始样本数量热力图...")
fig3, axes3 = plt.subplots(2, 3, figsize=(18, 10), sharex=True, sharey=True)
axes3 = axes3.flatten()

vmax_global = combined_df_all_types.groupby(['cbh_bin', 'thickness_bin'], observed=False).size().unstack().max().max() if not combined_df_all_types.empty else 1

for i, df in enumerate(all_data_loaded):
    if df.empty: continue
    ax = axes3[i]
    sample_counts = df.groupby(['cbh_bin', 'thickness_bin'], observed=False).size().unstack(fill_value=0)
    sns.heatmap(sample_counts, annot=True, fmt="d", cmap="Greens", linewidths=.5, ax=ax, cbar=False, norm=mcolors.LogNorm(vmin=1, vmax=vmax_global))
    ax.set_title(f'{group_names[i]}', fontsize=14, fontweight='bold')
    ax.set_ylabel('CloudSat CBH')
    ax.set_xlabel('Cloud Thickness')
    ax.invert_yaxis()

# 绘制 "All Types"
if not combined_df_all_types.empty:
    ax_all_count = axes3[-1]
    combined_counts = combined_df_all_types.groupby(['cbh_bin', 'thickness_bin'], observed=False).size().unstack(fill_value=0)
    sns.heatmap(combined_counts, annot=True, fmt="d", cmap="Greens", linewidths=.5, ax=ax_all_count, cbar=False, norm=mcolors.LogNorm(vmin=1, vmax=vmax_global))
    ax_all_count.set_title('All Types', fontsize=14, fontweight='bold')
    ax_all_count.set_ylabel('CloudSat CBH Bin')
    ax_all_count.set_xlabel('Cloud Thickness Bin')
    ax_all_count.invert_yaxis()

# 添加颜色条
cbar_ax_count = fig3.add_axes([0.92, 0.15, 0.02, 0.7])
norm_count = mcolors.LogNorm(vmin=1, vmax=vmax_global)
fig3.colorbar(plt.cm.ScalarMappable(cmap="Greens", norm=norm_count), cax=cbar_ax_count, label='Sample Count (Log Scale)')
# fig3.suptitle('图 3: 各云类型下原始样本数量分布', fontsize=20)
# fig3.tight_layout(rect=[0, 0, 0.9, 0.95])
plt.show()